# Copyright (c) 2020 Presto Labs Pte. Ltd.
# Author: jhkim

from absl import app, flags
from concurrent.futures import ProcessPoolExecutor

import copy
import functools
import logging

import numpy
import pandas
from pandas.plotting import register_matplotlib_converters
register_matplotlib_converters()
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

from matplotlib.backends.backend_pdf import PdfPages

from coin.base.book.types import TradeSide
from coin.base.datetime_util import to_timestamp_int

import coin.strategy.mm.tool.archive_base as abase


class BookDepth(object):
  def __init__(self, products, start_ts=None, end_ts=None, product_to_sub_request=None):
    self._products = set(products)
    self._start_ts = to_timestamp_int(start_ts or 0)
    self._end_ts = to_timestamp_int(end_ts or 0)
    self._bidsum = {}
    self._asksum = {}
    self._n = {}

  def on_book_reset(self, book_builder_name, book_builder):
    for product in self._products:
      book_builder.subscribe(product, functools.partial(self.on_book, product))
      if hasattr(book_builder, "subscribe_trade"):
        book_builder.subscribe_trade(product, functools.partial(self.on_trade, product))

  def on_book(self, product, book):
    if book.timestamp < self._start_ts:
      return
    if self._end_ts and self._end_ts < book.timestamp:
      return
    if product not in self._products:
      return

    askarr = book.get_ask_array(20)
    bidarr = book.get_bid_array(20)
    if len(askarr) == 0 or len(bidarr) == 0:
      return

    midp = 0.5 * (askarr[0, 0] + bidarr[0, 0])

    askarr[:, 0] = (((askarr[:, 0] - midp) / midp) * 1e5).astype(int)
    bidarr[:, 0] = (((-bidarr[:, 0] + midp) / midp) * 1e5).astype(int)

    asklevs = pandas.DataFrame(askarr).groupby(0).sum().cumsum()
    bidlevs = pandas.DataFrame(bidarr).groupby(0).sum().cumsum()

    try:
      askl = asklevs.interpolate('index').loc[numpy.linspace(0, 60,
                                                             61)].fillna(method='ffill').fillna(0)
      bidl = bidlevs.interpolate('index').loc[numpy.linspace(0, 60,
                                                             61)].fillna(method='ffill').fillna(0)
    except Exception:
      return

    if product not in self._bidsum:
      self._bidsum[product] = bidl
      self._asksum[product] = askl
      self._n[product] = 1
    else:
      self._bidsum[product] += bidl
      self._asksum[product] += askl
      self._n[product] += 1

  def on_trade(self, product, trade):
    return

  def plot_all(self):
    for product in self._bidsum.keys():
      bidarr = self._bidsum[product] / self._n[product]
      askarr = self._asksum[product] / self._n[product]
      bidarr = bidarr.reset_index().to_numpy()
      askarr = askarr.reset_index().to_numpy()
      plt.plot(bidarr[:, 0] / 10, bidarr[:, 1], 'g-')
      plt.plot(-askarr[:, 0] / 10, askarr[:, 1], 'r-')
      plt.title(product)


from coin.strategy.mm.tool.feed_archive_checker import (get_feedcache_machines)


def launch(baseflags, fsflags, symbols, out_dir):
  plt.rcParams['figure.figsize'] = 12, 8
  plt.rcParams['font.size'] = 9
  plt.rcParams['legend.fontsize'] = 9
  plt.rcParams['xtick.labelsize'] = 9
  plt.rcParams['ytick.labelsize'] = 9

  products = abase.symbols_to_products(symbols, baseflags)
  start_time, end_time = abase.get_time_range(baseflags)

  feedsub_request = abase.products_to_feedsub_request_with_api_override(
      products, fsflags.api_override, fsflags.feed_machine)

  product_to_sub_request = {}
  for sub_request, products in feedsub_request.gen_iter():
    for product in products:
      product_to_sub_request[product] = sub_request

  strategy = BookDepth(products=products,
                       start_ts=start_time,
                       end_ts=end_time,
                       product_to_sub_request=product_to_sub_request)

  abase.run_from_archive_base(baseflags=baseflags,
                              feed_sub_flags_proto=fsflags,
                              products=products,
                              on_book_reset=strategy.on_book_reset)

  return strategy


def main(argv):
  trading_date_given = flags.FLAGS.trading_date

  fsflags = abase.get_feed_subscription_flags()

  fs_overrides = [None]

  if flags.FLAGS.subsymbols is not None:
    fs_overrides = list(abase.get_flags_with_full_symbols(flags.FLAGS.subsymbols))

  out_dir = flags.FLAGS.out_dir
  trading_dates = abase.get_trading_dates(flags.FLAGS.trading_date)

  params = []

  time_range = flags.FLAGS.time_range

  for trading_date in trading_dates:
    flags.FLAGS.trading_date = trading_date.strftime("%Y%m%d")
    for trange in abase.split_time_range(
        time_range,
        flags.FLAGS.split_hours) if flags.FLAGS.split_hours is not None else [time_range]:
      cgroups = []
      for fs_override in fs_overrides:
        abase.set_fs_dict(fs_override)
        symbols = flags.FLAGS.symbols.split(",")

        baseflags = abase.get_base_flags()
        baseflags.time_range = trange
        param = [launch, copy.deepcopy(baseflags), copy.deepcopy(fsflags), list(symbols), out_dir]
        params.append(param)

  for param in params:
    (param[0](*param[1:])).plot_all()
    plt.savefig("pic/bookdepth.png")
    plt.close()
  else:
    futures = []
    outs = []
    with ProcessPoolExecutor(max_workers=8) as executor:
      for param in params:
        futures.append(executor.submit(param[0], *param[1:]))
      for future in futures:
        outs.append(future.result())
    filename = f"bookdepths"
    plt.rcParams['figure.figsize'] = 12, 8
    plt.rcParams['font.size'] = 9
    plt.rcParams['legend.fontsize'] = 9
    plt.rcParams['xtick.labelsize'] = 9
    plt.rcParams['ytick.labelsize'] = 9
    with PdfPages(f"pic/{filename}.pdf") as pdf:
      for out in outs:
        out.plot_all()
        pdf.savefig()
        plt.close()


if __name__ == '__main__':
  abase.define_base_flags()
  abase.define_feed_archive_flags()
  # fill up default settings
  flags.FLAGS.feed_machine = 'feed-05.ap-northeast-1.aws'
  flags.DEFINE_string('subsymbols', None, 'subscript symbol name')
  flags.DEFINE_string('symbols', None, 'symbol name')
  flags.DEFINE_string('out_dir', 'pic', 'Output directory.')
  flags.DEFINE_string('postfix', None, 'filename postfix')
  flags.DEFINE_string('excel_file', None, '')
  flags.DEFINE_float('split_hours', None, '')
  flags.DEFINE_bool('print_all_feed', False, '')
  flags.DEFINE_integer('num_cpu', 12, '')

  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')

  app.run(main)
